package data

import (
	"context"
	"fmt"
	"log"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// MongoDBHandler 处理与MongoDB的交互
type MongoDBHandler struct {
	client     *mongo.Client
	dbName     string
	collection string
}

// NewMongoDBHandler 创建一个新的MongoDB处理器
func NewMongoDBHandler(uri string, dbName string, collection string) (*MongoDBHandler, error) {
	// 设置客户端连接选项
	clientOptions := options.Client().ApplyURI(uri)

	// 连接到MongoDB
	client, err := mongo.Connect(context.Background(), clientOptions)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to mongodb: %v", err)
	}

	// 检查连接
	err = client.Ping(context.Background(), nil)
	if err != nil {
		client.Disconnect(context.Background())
		return nil, fmt.Errorf("failed to ping mongodb: %v", err)
	}

	log.Printf("Connected to MongoDB: %s, database: %s, collection: %s", uri, dbName, collection)

	return &MongoDBHandler{
		client:     client,
		dbName:     dbName,
		collection: collection,
	}, nil
}

// StoreDocument 存储文档到MongoDB
func (mh *MongoDBHandler) StoreDocument(document interface{}) error {
	// 获取集合
	collection := mh.client.Database(mh.dbName).Collection(mh.collection)

	// 插入文档
	_, err := collection.InsertOne(context.Background(), document)
	if err != nil {
		return fmt.Errorf("failed to insert document: %v", err)
	}

	log.Printf("Document stored successfully in MongoDB")
	return nil
}

// FindDocuments 根据查询条件查找文档
func (mh *MongoDBHandler) FindDocuments(filter interface{}, limit int64) ([]bson.M, error) {
	// 获取集合
	collection := mh.client.Database(mh.dbName).Collection(mh.collection)

	// 设置查询选项
	findOptions := options.Find()
	findOptions.SetLimit(limit)

	// 执行查询
	cursor, err := collection.Find(context.Background(), filter, findOptions)
	if err != nil {
		return nil, fmt.Errorf("failed to find documents: %v", err)
	}
	defer cursor.Close(context.Background())

	// 遍历结果
	var results []bson.M
	if err = cursor.All(context.Background(), &results); err != nil {
		return nil, fmt.Errorf("failed to decode documents: %v", err)
	}

	return results, nil
}

// GetDeviceData 获取指定设备的最新数据
func (mh *MongoDBHandler) GetDeviceData(deviceID string, limit int64) ([]bson.M, error) {
	// 创建过滤器，按设备ID查询
	filter := bson.M{"device_id": deviceID}

	// 设置排序选项，按时间戳降序排列
	sortOptions := options.Find()
	sortOptions.SetSort(bson.D{{Key: "timestamp", Value: -1}})
	sortOptions.SetLimit(limit)

	// 获取集合
	collection := mh.client.Database(mh.dbName).Collection(mh.collection)

	// 执行查询
	cursor, err := collection.Find(context.Background(), filter, sortOptions)
	if err != nil {
		return nil, fmt.Errorf("failed to find device data: %v", err)
	}
	defer cursor.Close(context.Background())

	// 遍历结果
	var results []bson.M
	if err = cursor.All(context.Background(), &results); err != nil {
		return nil, fmt.Errorf("failed to decode device data: %v", err)
	}

	return results, nil
}

// CreateIndexes 创建索引以提高查询性能
func (mh *MongoDBHandler) CreateIndexes() error {
	// 获取集合
	collection := mh.client.Database(mh.dbName).Collection(mh.collection)

	// 创建设备ID和时间戳的索引
	indexModels := []mongo.IndexModel{
		{
			Keys: bson.D{{Key: "device_id", Value: 1}},
		},
		{
			Keys: bson.D{{Key: "timestamp", Value: -1}},
		},
		{
			Keys: bson.D{{Key: "device_id", Value: 1}, {Key: "timestamp", Value: -1}},
		},
	}

	// 创建索引
	_, err := collection.Indexes().CreateMany(context.Background(), indexModels)
	if err != nil {
		return fmt.Errorf("failed to create indexes: %v", err)
	}

	log.Printf("Created indexes for MongoDB collection %s", mh.collection)
	return nil
}

// Close 关闭MongoDB连接
func (mh *MongoDBHandler) Close() error {
	if mh.client != nil {
		if err := mh.client.Disconnect(context.Background()); err != nil {
			return fmt.Errorf("failed to disconnect from mongodb: %v", err)
		}
		log.Println("Disconnected from MongoDB")
	}
	return nil
}
